import requests
import json

API_KEY = "2w2rFHt1hcfdRDE2kbzJs6Ks"
SECRET_KEY = "rnNScQ5TR0z9i3Boi3LzPZAuWSeDuyT0"

class WenXinService:

    # 发送机器人
    def send_msg():
        url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + ChatService.get_access_token()
        table_schema = {
            "tableName": "yy_user_member",
            "description": "用户表",
            "columns": [
                {"name": "id", "type": "INT", "description": "自增ID"},
                {"name": "member_uid", "type": "VARCHAR", "description": "用户唯一ID"},
                {"name": "full_name", "type": "VARCHAR", "description": "用户姓名"},
                {"name": "phone", "type": "VARCHAR", "description": "用户手机号"},
                {"name": "gender", "type": "INT", "description": "用户性别：1男，2女"},
                {"name": "created_at", "type": "datetime", "description": "用户注册时间"}
            ]
        }
        # content = "你好"
        content = "根据表结构，写出一条查询用户总数的sql语句，表结构为："+ json.dumps(table_schema)
        messages = [ 
            {"role": "user","content": content}
        ]
        payload = json.dumps({
            "messages" : messages,
            "temperature": 0.95,
            "top_p": 0.8,
            "penalty_score": 1,
            "enable_system_memory": False,
            "disable_search": False,
            "enable_citation": False,
            # boolean
            "response_format": {
                "type":"json", 
                "template":{
                    "greeting":"{{answer}}"
                }
            }
        })
        headers = {
            'Content-Type': 'application/json'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        print(response.text)
    

    def get_access_token():
        """
        使用 AK，SK 生成鉴权签名（Access Token）
        :return: access_token，或是None(如果错误)
        """
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
        return str(requests.post(url, params=params).json().get("access_token"))